import os
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from glob import glob
import json
from tqdm import tqdm
import pandas as pd
import h5py
from transformers import AutoTokenizer

import os

tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')

def find_npy_files(directory):
    """查找指定目录中的所有 NPY 文件"""
    npy_files = glob(os.path.join(directory, "**", "*.npy"), recursive=True)
    
    if not npy_files:
        print(f"在目录 {directory} 中未找到任何 .npy 文件")
        return []
    
    print(f"找到 {len(npy_files)} 个 NPY 文件:")
    for i, file in enumerate(npy_files[:5]):
        print(f"{i+1}. {file}")
    if len(npy_files) > 5:
        print(f"... 和另外 {len(npy_files)-5} 个文件")
    
    return npy_files

def analyze_npy_file(file_path):
    """分析单个 NPY 文件的内容"""
    print(f"\n分析文件: {file_path}")
    
    try:
        # 加载 NPY 文件
        data = np.memmap(file_path, dtype=np.uint16, mode = 'r')
        
        # 收集基本信息
        info = {
            "file_path": file_path,
            "file_size_mb": os.path.getsize(file_path) / (1024 * 1024),
            "data_type": str(type(data)),
            "shape": data.shape if hasattr(data, 'shape') else "N/A",
            "dtype": str(data.dtype) if hasattr(data, 'dtype') else "N/A"
        }
        
        # 打印基本信息
        print(f"文件大小: {info['file_size_mb']:.2f} MB")
        print(f"数据类型: {info['data_type']}")
        print(f"形状: {info['shape']}")
        print(f"数据类型: {info['dtype']}")
        
        # if isinstance(data, np.ndarray):
        #     print(f"值范围: {info['min']:.4f} 到 {info['max']:.4f}")
        #     print(f"平均值: {info['mean']:.4f} ± {info['std']:.4f}")
        
        return info, data
    
    except Exception as e:
        print(f"加载文件 {file_path} 失败: {str(e)}")
        return None, None

def visualize_npy_data(data, file_path):
    """可视化 NPY 文件中的数据"""
    if not isinstance(data, np.ndarray):
        print("无法可视化非数组数据")
        return
    
    print("\n创建数据可视化...")
    
    # 根据数据类型和维度选择合适的可视化方式
    if data.ndim == 1:
        plt.figure(figsize=(12, 6))
        plt.title(f"1D 数据分布: {os.path.basename(file_path)}")
        sns.histplot(data, bins=50, kde=True)
        plt.xlabel("值")
        plt.ylabel("频率")
        plt.savefig(f"{file_path}_1d_dist.png", bbox_inches='tight')
        plt.show()
        
    elif data.ndim == 2:
        # 对于大型矩阵，只显示一部分
        sample_size = min(1000, data.shape[0])
        sampled_data = data[:sample_size]
        
        plt.figure(figsize=(12, 8))
        plt.title(f"2D 数据热图 (前 {sample_size} 行): {os.path.basename(file_path)}")
        sns.heatmap(sampled_data, cmap="viridis", cbar_kws={'label': '数值'})
        plt.xlabel("列索引")
        plt.ylabel("行索引")
        plt.savefig(f"{file_path}_2d_heatmap.png", bbox_inches='tight')
        plt.show()
        
        # 值分布
        plt.figure(figsize=(12, 6))
        plt.title(f"2D 数据值分布: {os.path.basename(file_path)}")
        sns.histplot(data.flatten(), bins=50, kde=True)
        plt.xlabel("值")
        plt.ylabel("频率")
        plt.savefig(f"{file_path}_2d_dist.png", bbox_inches='tight')
        plt.show()
        
    elif data.ndim == 3:
        # 对于3D数据（如嵌入），显示切片
        slice_idx = min(10, data.shape[0] - 1)
        plt.figure(figsize=(12, 8))
        plt.title(f"3D 数据切片 (索引 {slice_idx}): {os.path.basename(file_path)}")
        sns.heatmap(data[slice_idx], cmap="viridis", cbar_kws={'label': '数值'})
        plt.xlabel("X 维度")
        plt.ylabel("Y 维度")
        plt.savefig(f"{file_path}_3d_slice.png", bbox_inches='tight')
        plt.show()
        
    else:
        print(f"不支持 {data.ndim} 维数据的可视化")

def find_related_metadata(npy_file):
    """查找与 NPY 文件相关的元数据文件"""
    base_path = os.path.splitext(npy_file)[0]
    possible_files = [
        f"{base_path}.json",
        f"{base_path}_metadata.json",
        f"{os.path.dirname(npy_file)}/metadata.json",
        f"{os.path.dirname(npy_file)}/config.json"
    ]
    
    for meta_file in possible_files:
        if os.path.exists(meta_file):
            return meta_file
    return None

def load_and_display_metadata(metadata_file):
    """加载并显示元数据文件内容"""
    print(f"\n找到关联的元数据文件: {metadata_file}")
    
    try:
        with open(metadata_file, 'r') as f:
            metadata = json.load(f)
            
        print("元数据内容:")
        print(json.dumps(metadata, indent=2, ensure_ascii=False))
        
        return metadata
    
    except Exception as e:
        print(f"加载元数据失败: {str(e)}")
        return None

def save_summary_report(file_infos, output_file="dolma_npy_report.html"):
    """生成 HTML 格式的总结报告"""
    print(f"\n生成总结报告: {output_file}")
    
    # 创建 DataFrame
    df = pd.DataFrame(file_infos)
    
    # 生成 HTML
    html_content = f"""
    <html>
    <head>
        <title>Dolma NPY 文件分析报告</title>
        <style>
            body {{ font-family: Arial, sans-serif; margin: 20px; }}
            h1 {{ color: #2c3e50; }}
            table {{ border-collapse: collapse; width: 100%; }}
            th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
            th {{ background-color: #f2f2f2; }}
            img {{ max-width: 100%; height: auto; margin: 10px 0; }}
        </style>
    </head>
    <body>
        <h1>Dolma NPY 文件分析报告</h1>
        <p>共分析 {len(file_infos)} 个文件</p>
        
        <h2>文件概览</h2>
        {df.to_html(index=False)}
        
        <h2>文件可视化</h2>
    """
    
    # 添加可视化图像
    for info in file_infos:
        base_name = os.path.basename(info['file_path'])
        image_files = glob(f"{info['file_path']}_*.png")
        
        if image_files:
            html_content += f"<h3>{base_name}</h3>"
            for img in image_files:
                img_name = os.path.basename(img)
                html_content += f"<div><h4>{img_name}</h4><img src='{img_name}'></div>"
    
    html_content += """
    </body>
    </html>
    """
    
    # 保存 HTML 文件
    with open(output_file, 'w') as f:
        f.write(html_content)
    
    print(f"报告已保存至: {output_file}")

def main():
    # 设置 Dolma 输出目录（根据您的实际情况修改）
    dolma_output_dir = "./python/dolma/data"  # 更改为您的实际路径
    
    if not os.path.exists(dolma_output_dir):
        print(f"目录不存在: {dolma_output_dir}")
        return
    
    # 查找所有 NPY 文件
    npy_files = find_npy_files(dolma_output_dir)
    if not npy_files:
        return
    
    all_file_infos = []
    
    # 分析每个文件
    for file_path in tqdm(npy_files, desc="分析 NPY 文件"):
        file_info, data = analyze_npy_file(file_path)
        
        if file_info is None:
            continue
            
        # 添加到汇总信息
        all_file_infos.append(file_info)
        
        # # 可视化数据
        # if data is not None:
        #     visualize_npy_data(data, file_path)
        data = np.memmap(file_path, dtype=np.uint16, mode = 'r')
        tokens = data[10000:11000].astype(np.int32)
        print(tokenizer.decode(tokens.tolist()))
        # 查找并显示相关元数据
    #     metadata_file = find_related_metadata(file_path)
    #     if metadata_file:
    #         load_and_display_metadata(metadata_file)
    
    # # 生成总结报告
    # if all_file_infos:
    #     save_summary_report(all_file_infos)

if __name__ == "__main__":
    main()